import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data


def weight_variable(shape):
    return tf.Variable(tf.random_normal(shape=shape, mean=0.0, stddev=1.0))

def bias_variable(shape):
    return tf.Variable(tf.constant(0.0, shape=shape))
def compute(y_true, y_predict):
    with tf.variable_scope("compute_loss"):
        loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=y_true, logits=y_predict))
    return loss

def model():
    # 准备占位符, x [None, 784], b: [None, 10]
    with tf.variable_scope('data'):
        x_data = tf.placeholder(tf.float32, [None, 784])
        y_true = tf.placeholder(tf.int32, [None, 10])
    # 卷积层1
    with tf.variable_scope('conv1'):
        # 权重的形状: [5, 5, 1, 32], b: [32] 宽高为5, 通道为1, 32个filter
        w_conv1 = weight_variable(shape=[5,5,1,32])
        b_conv1 = bias_variable([32]) # 32个偏置值
        # 进行卷积, relu激活, 池化操作
        # 将数据的形状处理成卷积层需要的数据格式
        # 将读取到的批量数据形状修改为, 数量待定, 宽高为28, 通道为1
        x_reshape = tf.reshape(x_data, [-1, 28, 28, 1])
        # 得到的形状 [None, 28, 28, 32]   [-1, 28, 28, 1] 通过 [5,5,1,32])filter 来扫描, 得到的形状是 [None,28,28,32]
        x_relu1 = tf.nn.relu(tf.nn.conv2d(x_reshape, w_conv1, strides=[1,1,1,1], padding="SAME") + b_conv1)
        # 经过池化之后是 [None, 14, 14, 32] # 池化操作, 减少特征
        x_pool1 = tf.nn.max_pool(x_relu1, ksize=[1,2,2,1], strides=[1,2,2,1], padding='SAME')
    # 卷积层2 64个filter,  大小5*5, 步长为1, padding为 "SAME"
    with tf.variable_scope('conv2'):
        # 上一层的输出 是下一层的输入 [None, 14, 14, 32]
        # 根据输入 准不权重和偏置
        w_conv2 = weight_variable(shape=[5,5,32,64])
        b_conv2 = bias_variable([64])
        # [None, 14, 14, 64]
        x_relu2 = tf.nn.relu(tf.nn.conv2d(x_pool1, w_conv2, strides=[1,1,1,1], padding='SAME') + b_conv2)
        # [None, 7, 7, 64]
        x_pool2 = tf.nn.max_pool(x_relu2, ksize=[1,2,2,1], strides=[1,2,2,1], padding='SAME')
    # conv2 输出的结果是 [None, 7, 7, 64] 全连接层 最终输出的结果是 [None, 10] -->[None,7,7,64]  * [7*7*64, 10] = [None, 10]
    with tf.variable_scope('fc'):
        w_fc = weight_variable(shape=[7*7*64, 10])
        b_fc = bias_variable([10])
        x_fc_reshape = tf.reshape(x_pool2, [-1, 7*7*64])
        y_predict = tf.matmul(x_fc_reshape, w_fc) + b_fc
    return x_data, y_true, y_predict

def sgd(loss, y_true, y_predict):
    with tf.variable_scope("SGD"):
        # 优化
        train_op = tf.train.GradientDescentOptimizer(0.001).minimize(loss)
    equal_list = tf.equal(tf.argmax(y_true, 1), tf.argmax(y_predict, 1))
    accuracy = tf.reduce_mean(tf.cast(equal_list, tf.float32))
    return  train_op, accuracy


def main(argv):
    mnist = input_data.read_data_sets('./data/mnist/input_data', one_hot=True)
    x_data, y_true, y_predict = model()
    loss = compute(y_true, y_predict)
    train_op, accuracy = sgd(loss, y_true, y_predict)

    init_op = tf.global_variables_initializer()

    with tf.Session() as sess:
        sess.run(init_op)
        for i in  range(2000):
            mnist_x, mnist_y = mnist.train.next_batch(50)
            if i % 100 == 0:
                print('训练的准确率: ', sess.run(accuracy, feed_dict={x_data: mnist_x, y_true: mnist_y}))
            sess.run(train_op, feed_dict={x_data: mnist_x, y_true: mnist_y})

if __name__ == '__main__':
    tf.app.run()
